package com.gitee.cirnochat.common.utils;

import com.gitee.cirnochat.chatai.model.dto.markdown.ImageCountResult;
import com.vladsch.flexmark.ast.Image;
import com.vladsch.flexmark.ext.media.tags.AudioLink;
import com.vladsch.flexmark.ext.media.tags.PictureLink;
import com.vladsch.flexmark.ext.media.tags.VideoLink;
import com.vladsch.flexmark.ext.media.tags.internal.AbstractMediaLink;
import com.vladsch.flexmark.html.HtmlRenderer;
import com.vladsch.flexmark.parser.Parser;
import com.vladsch.flexmark.profile.pegdown.Extensions;
import com.vladsch.flexmark.profile.pegdown.PegdownOptionsAdapter;
import com.vladsch.flexmark.util.ast.Node;
import com.vladsch.flexmark.util.ast.NodeVisitor;
import com.vladsch.flexmark.util.ast.TextCollectingVisitor;
import com.vladsch.flexmark.util.ast.VisitHandler;
import com.vladsch.flexmark.util.data.DataHolder;
import com.vladsch.flexmark.util.data.MutableDataSet;

import java.util.concurrent.atomic.AtomicInteger;

public final class MarkdownUtil {
    private MarkdownUtil() {
        // 私有构造函数，防止实例化
        throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
    }

    private final static DataHolder OPTIONS = PegdownOptionsAdapter.flexmarkOptions(Extensions.ALL);

    private static final MutableDataSet FORMAT_OPTIONS = new MutableDataSet();

    static {
        //copy extensions from Pegdown compatible to Formatting
        FORMAT_OPTIONS.set(Parser.EXTENSIONS, Parser.EXTENSIONS.get(OPTIONS));
    }

    private static final Parser parser = Parser.builder(OPTIONS).build();
    private static final HtmlRenderer render = HtmlRenderer.builder(OPTIONS).build();

    /**
     * 统计图片的数量
     *
     * @param markdown Markdown文本
     * @return 图片数量
     */
    public static ImageCountResult countImages(String markdown) {
        Node document = parser.parse(markdown);
        // 计数器
        AtomicInteger imageCount = new AtomicInteger();
        AtomicInteger mediaCount = new AtomicInteger();
        // 创建 AST 访问器，只访问 Image 节点
        NodeVisitor visitor = new NodeVisitor(
                new VisitHandler<>(Image.class, image -> {
                    String url = image.getUrl().toString().toLowerCase();
                    if (isImageUrl(url)) {
                        imageCount.getAndIncrement();
                    } else {
                        mediaCount.getAndIncrement();
                    }
                })
        );

        visitor.visit(document);
        return new ImageCountResult(imageCount.get(), mediaCount.get());
    }


    /**
     * 检查给定的URL是否是图片链接
     *
     * @param url URL字符串
     * @return 如果是图片链接则返回true，否则返回false
     */
    private static boolean isImageUrl(String url) {
        return url.matches(".*\\.(png|jpg|jpeg|gif|bmp|svg|webp)$");
    }

    /**
     * 统计Markdown中的媒体链接数量
     *
     * @param markdown Markdown文本
     * @return 媒体数量
     */
    public static int countMedia(String markdown) {
        Node document = parser.parse(markdown);
        // 计数器
        AtomicInteger imageCount = new AtomicInteger();
        // 创建 AST 访问器，只访问 Image 节点
        NodeVisitor visitor = new NodeVisitor(
                new VisitHandler<>(AbstractMediaLink.class, image -> {
                    imageCount.getAndIncrement();
                })
        );
        visitor.visit(document);
        return imageCount.get();
    }

    /**
     * 统计Markdown中的媒体图片链接数量
     *
     * @param markdown Markdown文本
     * @return 媒体数量
     */
    public static int countMediaPicture(String markdown) {
        Node document = parser.parse(markdown);
        // 计数器
        AtomicInteger imageCount = new AtomicInteger();
        // 创建 AST 访问器，只访问 Image 节点
        NodeVisitor visitor = new NodeVisitor(
                new VisitHandler<>(PictureLink.class, image -> {
                    imageCount.getAndIncrement();
                })
        );
        visitor.visit(document);
        return imageCount.get();
    }

    /**
     * 统计Markdown中的媒体音频链接数量
     *
     * @param markdown Markdown文本
     * @return 媒体数量
     */
    public static int countMediaAudio(String markdown) {
        Node document = parser.parse(markdown);
        // 计数器
        AtomicInteger audioCount = new AtomicInteger();
        // 创建 AST 访问器，只访问 Image 节点
        NodeVisitor visitor = new NodeVisitor(
                new VisitHandler<>(AudioLink.class, audio -> {
                    audioCount.getAndIncrement();
                })
        );
        visitor.visit(document);
        return audioCount.get();
    }


    /**
     * 统计Markdown中的媒体视频链接数量
     *
     * @param markdown Markdown文本
     * @return 媒体数量
     */
    public static int countMediaVideo(String markdown) {
        Node document = parser.parse(markdown);
        // 计数器
        AtomicInteger videoCount = new AtomicInteger();
        // 创建 AST 访问器，只访问 Image 节点
        NodeVisitor visitor = new NodeVisitor(
                new VisitHandler<>(VideoLink.class, video -> {
                    videoCount.getAndIncrement();
                })
        );
        visitor.visit(document);
        return videoCount.get();
    }

    /**
     * 提取Markdown中的纯文本信息
     *
     * @param markdown Markdown文本
     * @return 纯文本内容
     */
    public static String extractText(String markdown) {
        Node document = parser.parse(markdown);
        TextCollectingVisitor collectingVisitor = new TextCollectingVisitor();
        return collectingVisitor.collectAndGetText(document);
    }

    public static void main(String[] args) {
        String md = """
                # 标题
                这是一个Markdown文档。
                
                ![图片.png](https://example.com/image.png)
                
                ![音频.mp3](https://example.com/audio.mp3)
                
                ![视频.mp4](https://example.com/video.mp4)
                
                <audio src="https://example.com/audio.mp3" controls></audio>
                
                <video src="https://example.com/video.mp4" controls></video>
                """;

        System.out.println("图片数量: " + countImages(md));
        System.out.println("媒体链接数量: " + countMedia(md));
        System.out.println("媒体图片链接数量: " + countMediaPicture(md));
        System.out.println("媒体音频链接数量: " + countMediaAudio(md));
        System.out.println("媒体视频链接数量: " + countMediaVideo(md));
        System.out.println("纯文本内容: " + extractText(md));
    }

}
